package com.zoe.flume.sink.elasticsearch;

import org.apache.flume.Event;
import org.apache.flume.EventDeliveryException;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.InetSocketTransportAddress;
import org.elasticsearch.transport.client.PreBuiltTransportClient;

import java.net.InetAddress;
import java.net.UnknownHostException;

import static com.zoe.flume.sink.elasticsearch.ElasticSearchSinkConstants.DEFAULT_PORT;

/**
 * Created by shenyuhang on 2018/2/1.
 */
public class ElasticSearchTransportClient {

    private TransportClient transportClient;
    private BulkRequestBuilder bulkRequestBuilder;
    private IndexRequestBuilder indexRequestBuilder;
    private ElasticSearchEventSerializer serializer;
    private String indexName;
    private String indexType;


    private static class Builder {
        private final static ElasticSearchTransportClient client = new ElasticSearchTransportClient();

        public static ElasticSearchTransportClient get() {
            return client;
        }
    }

    public static ElasticSearchTransportClient builder() {
        return Builder.get();
    }

    public ElasticSearchTransportClient setTransportClient(String clusterName, String[] esHostAndPorts) {
        Settings settings = Settings.builder().put("cluster.name", clusterName).build();

        for (String esHostAndPort : esHostAndPorts) {
            String[] hostAndPort = new String[2];
            String[] hostPort = esHostAndPort.trim().split(":");
            hostAndPort[0] = hostPort[0];
            if (hostPort.length < 2) {
                hostAndPort[1] = Integer.toString(DEFAULT_PORT);
            } else {
                hostAndPort[1] = hostPort[1];
            }

            try {
                if (null == transportClient) {
                    transportClient = new PreBuiltTransportClient(settings)
                            .addTransportAddress(new InetSocketTransportAddress(InetAddress.getByName(hostAndPort[0]), Integer.parseInt(hostAndPort[1])));
                } else {
                    transportClient = transportClient.addTransportAddress(new InetSocketTransportAddress(InetAddress.getByName(hostAndPort[0]), Integer.parseInt(hostAndPort[1])));
                }
            } catch (UnknownHostException e) {
                e.printStackTrace();
            }
        }
        return this;
    }

    public ElasticSearchTransportClient setBulk() {
        bulkRequestBuilder = transportClient.prepareBulk();
        return this;
    }

    public BulkRequestBuilder getBulkRequestBuilder() {
        return bulkRequestBuilder;
    }

    public ElasticSearchTransportClient setIndexName(String indexName) {
        this.indexName = indexName;
        return this;
    }

    public ElasticSearchTransportClient setIndexType(String indexType) {
        this.indexType = indexType;
        return this;
    }

    public IndexRequestBuilder getIndexRequestBuilder() {
        return transportClient.prepareIndex(indexName, indexType);
    }

    public ElasticSearchTransportClient setSerializer(ElasticSearchEventSerializer serializer) {
        this.serializer = serializer;
        return this;
    }


    public void addEvent(Event event) throws Exception {
        bulkRequestBuilder.add(getIndexRequestBuilder().setSource(serializer.serializer(event)));
    }

    public void submit() throws Exception {
        try {
            BulkResponse bulkResponse = bulkRequestBuilder.execute().actionGet();
            if (bulkResponse.hasFailures()) {
                throw new EventDeliveryException(bulkResponse.buildFailureMessage());
            }
        } finally {
            bulkRequestBuilder = transportClient.prepareBulk();
        }
    }

    public void close() {
        if (transportClient != null) {
            transportClient.close();
        }
        transportClient = null;
        bulkRequestBuilder = null;
        indexRequestBuilder = null;
    }


}
